package lexer;

import java.util.HashMap;
import java.util.Map;

public class Lexer {
    private String input;
    private int position;
    private int line;
    private int column;
    private char currentChar;

    private static final Map<String, TokenType> KEYWORDS = new HashMap<>();

    static {
        // 关键字
        KEYWORDS.put("IF", TokenType.IF);
        KEYWORDS.put("THEN", TokenType.THEN);
        KEYWORDS.put("ELSE", TokenType.ELSE);
        KEYWORDS.put("ENDIF", TokenType.ENDIF);
        KEYWORDS.put("FOR", TokenType.FOR);
        KEYWORDS.put("TO", TokenType.TO);
        KEYWORDS.put("STEP", TokenType.STEP);
        KEYWORDS.put("DO", TokenType.DOO);
        KEYWORDS.put("ENDFOR", TokenType.ENDFOR);
        KEYWORDS.put("WHILE", TokenType.WHILE);
        KEYWORDS.put("ENDWHILE", TokenType.ENDWHILE);
        KEYWORDS.put("BREAK", TokenType.BREAK);
        KEYWORDS.put("CONTINUE", TokenType.CONTINUE);
        KEYWORDS.put("AND", TokenType.AND);
        KEYWORDS.put("OR", TokenType.OR);
        KEYWORDS.put("NOT", TokenType.NOT);
    }

    public Lexer(String input) {
        this.input = input;
        this.position = 0;
        this.line = 1;
        this.column = 0;
        this.currentChar = input.length() > 0 ? input.charAt(0) : '\0';
    }

    private void advance() {
        if (currentChar == '\n') {
            line++;
            column = 0;
        } else {
            column++;
        }

        position++;
        if (position < input.length()) {
            currentChar = input.charAt(position);
        } else {
            currentChar = '\0';
        }
    }

    private char peek() {
        if (position + 1 < input.length()) {
            return input.charAt(position + 1);
        }
        return '\0';
    }

    private void skipWhitespace() {
        while (currentChar != '\0' && Character.isWhitespace(currentChar)) {
            advance();
        }
    }

    private void skipComment() {
        while (currentChar != '\0' && currentChar != '\n') {
            advance();
        }
        if (currentChar == '\n') {
            advance();
        }
    }

    private Token number() {
        StringBuilder result = new StringBuilder();
        while (currentChar != '\0' && (Character.isDigit(currentChar) || currentChar == '.')) {
            result.append(currentChar);
            advance();
        }
        return new Token(TokenType.NUMBER, result.toString(), line, column);
    }

    private Token identifier() {
        StringBuilder result = new StringBuilder();
        while (currentChar != '\0' &&
                (Character.isLetterOrDigit(currentChar) || currentChar == '_')) {
            result.append(currentChar);
            advance();
        }

        String id = result.toString().toUpperCase();
        TokenType type = KEYWORDS.getOrDefault(id, TokenType.IDENTIFIER);
        return new Token(type, id, line, column);
    }

    private Token string() {
        StringBuilder result = new StringBuilder();
        advance(); // 跳过开始的引号

        while (currentChar != '\0' && currentChar != '"') {
            if (currentChar == '\\') {
                advance();
                switch (currentChar) {
                    case 'n': result.append('\n'); break;
                    case 't': result.append('\t'); break;
                    case '"': result.append('"'); break;
                    case '\\': result.append('\\'); break;
                    default: result.append('\\').append(currentChar); break;
                }
            } else {
                result.append(currentChar);
            }
            advance();
        }

        if (currentChar == '"') {
            advance();
        }

        return new Token(TokenType.STRING, result.toString(), line, column);
    }

    public Token getNextToken() {
        while (currentChar != '\0') {
            if (Character.isWhitespace(currentChar)) {
                skipWhitespace();
                continue;
            }

            if (currentChar == '/' && peek() == '/') {
                skipComment();
                continue;
            }

            if (Character.isDigit(currentChar)) {
                return number();
            }

            if (Character.isLetter(currentChar) || currentChar == '_') {
                return identifier();
            }

            if (currentChar == '"') {
                return string();
            }

            // 处理多字符运算符
            if (currentChar == ':' && peek() == '=') {
                advance();
                advance();
                return new Token(TokenType.ASSIGN, ":=", line, column);
            }

            if (currentChar == '>' && peek() == '=') {
                advance();
                advance();
                return new Token(TokenType.GE, ">=", line, column);
            }

            if (currentChar == '<' && peek() == '=') {
                advance();
                advance();
                return new Token(TokenType.LE, "<=", line, column);
            }

            if (currentChar == '!' && peek() == '=') {
                advance();
                advance();
                return new Token(TokenType.NE, "!=", line, column);
            }

            if (currentChar == '=') {
                advance();
                return new Token(TokenType.EQ, "=", line, column);
            }

            // 单字符token
            switch (currentChar) {
                case '+': advance(); return new Token(TokenType.PLUS, "+", line, column);
                case '-': advance(); return new Token(TokenType.MINUS, "-", line, column);
                case '*': advance(); return new Token(TokenType.MULTIPLY, "*", line, column);
                case '/': advance(); return new Token(TokenType.DIVIDE, "/", line, column);
                case '%': advance(); return new Token(TokenType.MOD, "%", line, column);
                case '>': advance(); return new Token(TokenType.GT, ">", line, column);
                case '<': advance(); return new Token(TokenType.LT, "<", line, column);
                case '(': advance(); return new Token(TokenType.LPAREN, "(", line, column);
                case ')': advance(); return new Token(TokenType.RPAREN, ")", line, column);
                case ',': advance(); return new Token(TokenType.COMMA, ",", line, column);
                case ';': advance(); return new Token(TokenType.SEMICOLON, ";", line, column);
                default:
                    Token error = new Token(TokenType.ERROR, String.valueOf(currentChar), line, column);
                    advance();
                    return error;
            }
        }

        return new Token(TokenType.EOF, "", line, column);
    }


    // 添加getter方法
    public String getInput() {
        return input;
    }

    public int getPosition() {
        return position;
    }

    public int getLine() {
        return line;
    }

    public int getColumn() {
        return column;
    }
}